package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"github.com/redis/go-redis/v9"
	_ "github.com/taosdata/driver-go/v3/taosWS"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Payload struct {
	Timestamp   string
	Tag_id      string
	Description string
	Unit        string
	Value       string
}

func (this *Payload) ToSQL() string {
	return "'" + this.Tag_id + "'" + " USING meters TAGS (" + "'" + this.Tag_id + "'" + ") VALUES (" + "'" + this.Timestamp + "'" + "," + this.Value + ")"
}

type Config struct {
	WorkerCount        int    `json:"workerCount"`
	BatchSize          int    `json:"batchSize"`
	TaosMaxOpenConns   int    `json:"taosMaxOpenConns"`
	TaosMaxIdleConns   int    `json:"taosMaxIdleConns"`
	TaosHost           string `json:"taosHost"`
	TaosPort           int    `json:"taosPort"`
	TaosUser           string `json:"taosUser"`
	TaosPassword       string `json:"taosPassword"`
	TaosDBName         string `json:"taosDBName"`
	RedisHost          string `json:"redisHost"`
	RedisPort          int    `json:"redisPort"`
	RedisDB            int    `json:"redisDB"`
	RedisPassword      string `json:"redisPassword"`
	RedisPoolSize      int    `json:"redisPoolSize"`
	RedisMinIdleConns  int    `json:"redisMinIdleConns"`
	RedisChannelBuffer int    `json:"redisChannelBuffer"`
	RedisListName      string `json:"redisListName"`
}

func loadConfig(file string) (*Config, error) {
	f, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var cfg Config
	if err := json.NewDecoder(f).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

var ctx = context.Background()

func ConnRedis(config *Config) (rdb *redis.Client) {
	rdb = redis.NewClient(&redis.Options{
		Addr:         config.RedisHost + ":" + strconv.Itoa(config.RedisPort),
		Password:     config.RedisPassword,
		DB:           config.RedisDB,
		PoolSize:     config.RedisPoolSize,
		MinIdleConns: config.RedisMinIdleConns,
	})
	err := rdb.Ping(ctx).Err()
	if err != nil {
		log.Println("链接失败", err)
		return
	}
	log.Println("链接成功 Redis")
	return rdb
}

func ConsumerRedisList(config *Config, realCh *chan Payload, redisCount *int64) {
	rdb := ConnRedis(config)
	for {
		result, err := rdb.BRPop(ctx, 0*time.Second, config.RedisListName).Result()
		if err != nil {
			log.Println("消费数据错误", err)
			time.Sleep(1 * time.Second)
			continue
		}
		if len(result) >= 2 {
			var data Payload
			err := json.Unmarshal([]byte(result[1]), &data)
			if err != nil {
				log.Println("序列化失败", err)
				time.Sleep(1 * time.Second)
				continue
			}
			// 更新 Redis 读取统计计数（每取到一条数据加1）
			atomic.AddInt64(redisCount, 1)
			*realCh <- data
		}
	}

}

func SetTaosConn(config *Config, taosDSN *string) (taos *sql.DB, err error) {
	// 无限重连直到成功建立连接
	for {
		taos, err = sql.Open("taosWS", *taosDSN)
		if err != nil {
			log.Println("Failed to connect to " + *taosDSN + "; ErrMessage: " + err.Error())
			time.Sleep(1 * time.Second)
			continue
		} else {
			taos.SetMaxOpenConns(config.TaosMaxOpenConns)
			taos.SetMaxIdleConns(config.TaosMaxIdleConns)
			// Ping 测试确保连接真正可用
			err = taos.Ping()
			if err != nil {
				log.Printf("数据库 Ping 失败: %v", err)
				time.Sleep(1 * time.Second)
			} else {
				log.Println("成功连接到 TDengine!")
				_, err := taos.Exec("create database if not exists archive vgroups 10 buffer 10;")
				if err != nil {
					log.Println("初始化失败", err)
					return nil, err
				}
				_, err = taos.Exec("CREATE STABLE IF NOT EXISTS archive.meters (ts TIMESTAMP, val VARCHAR(128)) TAGS (tag_id VARCHAR(128));")
				if err != nil {
					log.Println("创建超级表失败", err)
					return nil, err
				}
				return taos, nil
				// 成功连接，退出重连循环
			}
		}
		// 等待 1 秒后重连
		time.Sleep(1 * time.Second)
	}
}

func writeBatch(taos *sql.DB, batch []string, retries int) error {
	// 写入数据到 TDengine，调用写入接口
	var err error
	for attempt := 1; attempt <= retries; attempt++ {
		// 调用 TDengine 的无模式写入接口（此处采用 influxdb 格式，依据实际情况选择协议）
		batchString := "INSERT INTO " + strings.Join(batch, " ")
		_, err := taos.Exec("use archive;")
		if err != nil {
			log.Println("切换数据库失败")
			return err
		}
		_, err = taos.Exec(batchString)
		if err == nil {
			return nil
		}
		log.Printf("写入失败，尝试 %d/%d: %v", attempt, retries, err)
		time.Sleep(time.Second) // 重试前等待 1 秒（可根据需要调整等待时间）
	}
	return err
}

func InsertDatabase(config *Config, taos *sql.DB, realCh *chan Payload, tdEngineCount *int64) {
	// 启动 worker pool 协程 批量从 realCh 中读取数据并写入 tdengine
	var wg sync.WaitGroup
	for i := 0; i < config.WorkerCount; i++ {
		wg.Add(1)
		go func(workerId int) {
			defer wg.Done()
			// 每个协程独立维护自己的批次数据
			batch := make([]string, 0, config.BatchSize)
			// 使用 ticker 定时提交未满批次的数据，避免数据滞留
			ticker := time.NewTicker(1 * time.Second)
			defer ticker.Stop()
			for {
				select {
				case data, ok := <-*realCh:
					// 当通道关闭时，写入剩余数据退出
					if !ok {
						if len(batch) > 0 {
							if err := writeBatch(taos, batch, 3); err != nil {
								log.Printf("Worker %d 批次写入失败: %v", workerId, err)
							} else {
								atomic.AddInt64(tdEngineCount, int64(len(batch)))
								//log.Printf("Worker %d 最后批次写入 %d 条记录成功", workerId, len(batch))
							}
						}
						return
					}
					batch = append(batch, data.ToSQL())
					// 达到批量大小则写入 TDengine
					if len(batch) >= config.BatchSize {
						if err := writeBatch(taos, batch, 3); err != nil {
							log.Printf("Worker %d 批次写入失败: %v", workerId, err)
						} else {
							atomic.AddInt64(tdEngineCount, int64(len(batch)))
							//log.Printf("Worker %d 批次写入 %d 条记录成功", workerId, len(batch))
						}
						batch = batch[:0]
					}
				case <-ticker.C:
					// 定时检查并写入未满批的数据
					if len(batch) > 0 {
						if err := writeBatch(taos, batch, 3); err != nil {
							log.Printf("Worker %d 定时批次写入失败: %v", workerId, err)
						} else {
							atomic.AddInt64(tdEngineCount, int64(len(batch)))
							//log.Printf("Worker %d 定时批次写入 %d 条记录成功", workerId, len(batch))
						}
						batch = batch[:0]
					}
				}
			}
		}(i)
	}
}

// monitorMetrics 每分钟打印一次从 Redis 读取与 TDengine 写入的条数统计，并复位计数器
func monitorMetrics(redisCount *int64, tdEngineCount *int64) {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()
	for {
		<-ticker.C
		redisNum := atomic.SwapInt64(redisCount, 0)
		tdengineNum := atomic.SwapInt64(tdEngineCount, 0)
		log.Printf("【监控】过去1分钟：Redis读取 %d 条数据, TDengine写入 %d 条记录", redisNum, tdengineNum)
	}
}

func main() {
	// 打开或创建日志文件，设置为追加模式
	logFile, err := os.OpenFile("app.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		// 如果打开文件失败，则使用 log.Fatalf 打印错误并终止程序
		log.Fatalf("无法打开日志文件: %v", err)
	}
	// 在程序退出时关闭文件
	defer logFile.Close()
	// 设置日志输出到文件
	log.SetOutput(logFile)
	log.SetOutput(os.Stdout)
	// 设置日志格式（可选）
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	/*开始加载配置文件*/
	cfg, err := loadConfig("config.json")
	if err != nil {
		log.Println("加载配置文件失败", err)
		return
	}
	/*加载配置文件完成*/
	// redis 的 实时数据 channel 带缓冲
	var realCh = make(chan Payload, cfg.RedisChannelBuffer)
	// 创建一个全局的 taos 数据库链接
	// 首先初始化taos
	taosDSN := cfg.TaosUser + ":" + cfg.TaosPassword + "@ws(" + cfg.TaosHost + ":" + strconv.Itoa(cfg.TaosPort) + ")/"
	taos, err := SetTaosConn(cfg, &taosDSN)
	if err != nil {
		log.Println("初始化失败", err)
		return
	}
	var (
		redisCount    int64
		tdEngineCount int64
	)
	go ConsumerRedisList(cfg, &realCh, &redisCount)
	go InsertDatabase(cfg, taos, &realCh, &tdEngineCount)
	go monitorMetrics(&redisCount, &tdEngineCount)
	select {}
}
